#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：OA-DG-main 
@File    ：thyroid_voc.py
@IDE     ：PyCharm 
@Author  ：cao xu
@Date    ：2025/9/16 下午2:36 
输出：“dataset type: ThyroidVOCDataset
num samples: 66756
first img: /data/lining/data/Structured_Dataset/Thyroid_Data/Comprehensive_data/picture/images/上海十院/SHSY-00000001/SHSY-00000001-US01-N1-IMG01.jpg
first xml: /data/caoxu/dataset/div-align-dg/voc_labels/上海十院/SHSY-00000001/SHSY-00000001-US01-N1-IMG01.xml
keys: dict_keys(['img_metas', 'img', 'img2', 'gt_bboxes', 'gt_bboxes2', 'gt_labels', 'multilevel_boxes', 'oamix_boxes'])”
"""
# mmdet/datasets/thyroid_voc.py
import os.path as osp
import numpy as np
import cv2
import xml.etree.ElementTree as ET

from mmdet.datasets import DATASETS
from mmdet.datasets.voc import VOCDataset

@DATASETS.register_module()
class ThyroidVOCDataset(VOCDataset):
    """Read train/test .txt where each line is a relative image path.
    Map it to IMG_ROOT/<rel> and XML_ROOT/<rel>.xml, then parse VOC xml.
    """
    def __init__(self, img_root, xml_root, ann_file, **kwargs):
        self.img_root = img_root
        self.xml_root = xml_root
        # 设成空串，允许 filename 保存绝对路径
        super().__init__(ann_file=ann_file, img_prefix='VOC2007', **kwargs)

    def load_annotations(self, ann_file):
        data_infos = []
        with open(ann_file, 'r', encoding='utf-8') as f:
            for line in f:
                rel = line.strip()
                if not rel:
                    continue
                img_path = osp.join(self.img_root, rel)
                xml_path = osp.join(self.xml_root, rel.rsplit('.', 1)[0] + '.xml')

                # 读出宽高（优先 XML，缺失则回退到实际图像）
                w = h = None
                if osp.exists(xml_path):
                    try:
                        root = ET.parse(xml_path).getroot()
                        size = root.find('size')
                        if size is not None:
                            w = int(float(size.find('width').text))
                            h = int(float(size.find('height').text))
                    except Exception:
                        pass
                if w is None or h is None:
                    try:
                        img = cv2.imread(img_path)
                        if img is not None:
                            h, w = img.shape[:2]
                    except Exception:
                        pass
                if w is None or h is None:
                    raise FileNotFoundError(f'Cannot get size for {img_path} / {xml_path}')

                img_id = rel.replace('/', '_').rsplit('.', 1)[0]
                data_infos.append(dict(
                    id=img_id,
                    filename=img_path,     # 绝对路径
                    width=w,
                    height=h,
                    ann=dict(ann_file=xml_path)  # 保存真实 xml 路径，后续直接用
                ))
        self.data_infos = data_infos
        return data_infos

    def _filter_imgs(self, min_size=32):
        """覆盖默认的 VOC 过滤逻辑，使用我们记录的真实 xml 路径。"""
        valid_inds = []
        for i, info in enumerate(self.data_infos):
            img_ok = osp.exists(info['filename'])
            xml_ok = ('ann' in info and 'ann_file' in info['ann']
                      and osp.exists(info['ann']['ann_file']))
            if not (img_ok and xml_ok):
                continue
            if min(info.get('width', 0), info.get('height', 0)) >= min_size:
                valid_inds.append(i)
        return valid_inds

    def get_ann_info(self, idx):
        xml_path = self.data_infos[idx]['ann']['ann_file']
        return self._parse_xml(xml_path)

    def _parse_xml(self, xml_path):
        root = ET.parse(xml_path).getroot()
        bboxes, labels = [], []
        bboxes_ignore, labels_ignore = [], []
        for obj in root.findall('object'):
            name_node = obj.find('name')
            if name_node is None:
                continue
            name = name_node.text
            # 忽略不在 CLASSES 的标签，避免类别名不匹配报错
            if name not in self.CLASSES:
                continue
            difficult = 0
            diff_node = obj.find('difficult')
            if diff_node is not None and diff_node.text is not None:
                try:
                    difficult = int(float(diff_node.text))
                except Exception:
                    difficult = 0
            bnd = obj.find('bndbox')
            xmin = float(bnd.find('xmin').text)
            ymin = float(bnd.find('ymin').text)
            xmax = float(bnd.find('xmax').text)
            ymax = float(bnd.find('ymax').text)
            bbox = [xmin, ymin, xmax, ymax]
            label = self.cat2label[name]
            if difficult:
                bboxes_ignore.append(bbox)
                labels_ignore.append(label)
            else:
                bboxes.append(bbox)
                labels.append(label)

        def to_np(a, dtype, shape4=False):
            if len(a) == 0:
                return np.zeros((0, 4), dtype=dtype) if shape4 else np.zeros((0,), dtype=dtype)
            return np.array(a, dtype=dtype)

        return dict(
            bboxes=to_np(bboxes, np.float32, True),
            labels=to_np(labels, np.int64, False),
            bboxes_ignore=to_np(bboxes_ignore, np.float32, True),
            labels_ignore=to_np(labels_ignore, np.int64, False),
        )

    def get_cat_ids(self, idx):
        ann = self.get_ann_info(idx)
        return list(set(ann['labels'].tolist()))
